import numpy as np
import torch
import torchio
from torch.utils.data import Dataset

from utils.DataUtils import get_all_data
from utils.DataUtils import preprocess
from utils.LogUtil import my_logger


class KiPA2022Dataset(Dataset):
    def __init__(self, sample_path_list, configure, dataset_type="train", ):
        self.dataset_type = dataset_type
        self.configure = configure
        my_logger.info("The " + self.dataset_type + " data's size is " + str(len(sample_path_list)))
        self.img_list, self.label_list = get_all_data(sample_path_list, configure.patch_size, configure.crop_method,
                                                      configure.dataset_name, configure.run_type, configure.if_mask)
        my_logger.info("The loading of " + self.dataset_type + " data has been completed!")
        self.normal_transform = torchio.ZNormalization()

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        img = self.img_list[idx]
        label = self.label_list[idx]
        img = torch.from_numpy(np.expand_dims(img, axis=0).astype(np.int16))
        img = self.normal_transform(img)
        label = torch.from_numpy(label).long()
        if self.configure.class_num > 1:
            label = torch.nn.functional.one_hot(label, num_classes=self.configure.class_num)
            label = label.permute(3, 0, 1, 2)
        return img, label


class CustomReadDataset(Dataset):
    def __init__(self, sample_path_list, configure, dataset_type="train", ):
        self.dataset_type = dataset_type
        my_logger.info("The " + self.dataset_type + " data's size is " + str(len(sample_path_list)))
        self.img_list, self.label_list = get_all_data(sample_path_list, configure.patch_size, configure.crop_method,
                                                      configure.if_mask)
        my_logger.info("The loading of " + self.dataset_type + " data has been completed!")

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        img, label = self.img_list[idx], self.label_list[idx]
        img = preprocess(np.expand_dims(img, axis=0))
        label = np.expand_dims(label, axis=0)
        return img, label


class CustomDataset(Dataset):
    def __init__(self, img_list, label_list, dataset_type="train", ):
        self.img_list = img_list
        self.label_list = label_list
        self.dataset_type = dataset_type
        my_logger.info("The " + self.dataset_type + " data's size is " + str(len(self.img_list)))

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        img, label = self.img_list[idx], self.label_list[idx]

        # img = torch.from_numpy(img)
        # label = torch.from_numpy(label).long()
        # label = torch.nn.functional.one_hot(label)  # The shape is [W,H,L,C]

        return img, label
